# 要执行器沿着轨迹运动
# Implement a second loop to servo the basis on the line.
# 此程序通过嵌套循环，跟踪轨迹上的每个点来实现IK
# 这里参考一下循环中如何设定运动速度
import pinocchio as pin
from pinocchio.utils import rand, zero
import meshcat
import numpy as np
# 引入src中所有的函数
import sys
from os.path import dirname, join

sys.path.append(join(dirname(dirname(__file__)), "src")) # src的函数需要import src下的其他module
sys.path.append(dirname(dirname(__file__))) # 把文件根路径添加到sys.path中，这样在根路径下的src文件夹才能被import
from src import *
from EasyRobo import *
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# 获取URDF模型路径
URDF_FILE_PATH = filePath(URDF_NAME = 'ur5_gripper', ROBOT = True)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# 构建robot对象，并得到viz
robot, viz = robot_viz(URDF_FILE_PATH)
q0 = pin.randomConfiguration(robot.model)
pin.framesForwardKinematics(robot.model, robot.data, q0)
pin.updateFramePlacements(robot.model, robot.data) 
robot.initViewer(loadModel=True)
viz.display(q0)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
import time
import math
# 需要求伪逆
from numpy.linalg import pinv
# 设置初始位置
q = q0
idx = 6
print('idx is ',idx)
dt_ik = 0.01
ik_time_max = 50.0
IT_MAX = math.floor(ik_time_max / dt_ik)
damp = 1e-12
##########################################


# 设置effector目标位置和姿态，并进行IK求解，注意速度要保持均匀

t = 0.0
dt = 0.01
qlist = []
ypos_v = 0.5
ypos_des = 0.3

it_num_max = math.floor(ypos_des / (ypos_v * dt))

for it_num in range(it_num_max):
    tic = time.time() # 计时开始
    ypos = it_num*ypos_v*dt
    
    T_t = pin.SE3(np.eye(3), np.array([0.5, ypos, 0.0])) # 确定目标位置和姿态
    homo_matrix = T_t.homogeneous
    ############################################################
    # 创建一个box位于目标位置，方便查看
    box_name = "world/box"
    box_size = [0.04, 0.02, 0.01]

    viz.viewer[box_name].set_object(
        meshcat.geometry.Box(box_size),
        meshcat.geometry.MeshLambertMaterial(color=0x6622cc,reflectivity=0.4)
    )

    viz.viewer[box_name].set_transform(
        homo_matrix
    )


    # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    # 开始进行逆运动学求解
    i = 0
    
    while True:
        T_bt = robot.data.oMi[idx].actInv(T_t)
        J_b = pin.computeJointJacobian(robot.model, robot.data, q, idx)
        J_l = -pin.Jlog6(T_bt.inverse())

        # 接下来会这里会利用J_error v^{*} = -e(q)来求v^{*}
        J_error = J_l.dot(J_b)
        e_q = pin.log(T_bt).vector # 定义的误差向量

        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        v_star = pinv(J_error).dot(-e_q) # 此处是求解的关键，这里直接用pinv，没有借用damping项
        # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
        q = pin.integrate(robot.model, q, v_star*dt_ik)
        
        pin.forwardKinematics(robot.model, robot.data, q) # 更新位置，这个一定不能忘


        if np.linalg.norm(e_q) < 1e-4:
            success = True
            break
        if i>IT_MAX:
            success = False
            break
        i += 1

    qlist.append(q)# 记录每次收敛得到的q
    if success:
        print("Convergence achieved!")
    else:
        print("Convergence failed!")

    
    toc = time.time()
    elapsed = toc - tic
    dt_sleep = max(0, dt - (elapsed))
    time.sleep(dt_sleep)

# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
vqmax = 0.8
dt_inter_pos = 0.01
qlist_inter = InterPos(q0, qlist, vqmax, dt_inter_pos)


for q in qlist_inter:
    tic = time.time() # 计时开始
    pin.forwardKinematics(robot.model, robot.data, q) # 更新位置，这个一定不能忘
    viz.display(q)
    toc = time.time()
    elapsed = toc - tic
    dt_sleep = max(0, dt_inter_pos - (elapsed))
    time.sleep(dt_sleep)
    
